package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/jmoiron/sqlx"
)

type Config struct {
	Database struct {
		Host     string `json:"host"`
		Port     int    `json:"port"`
		Username string `json:"username"`
		Password string `json:"password"`
		DBName   string `json:"dbname"`
	} `json:"database"`
}

type StockData struct {
	Symbol       string    `db:"symbol"`
	Exchange     string    `db:"exchange"`
	DateTime     time.Time `db:"datetime"`
	Interval     string    `db:"interval"`
	Volume       float64   `db:"volume"`
	Turnover     float64   `db:"turnover"`
	OpenInterest float64   `db:"open_interest"`
	OpenPrice    float64   `db:"open_price"`
	HighPrice    float64   `db:"high_price"`
	LowPrice     float64   `db:"low_price"`
	ClosePrice   float64   `db:"close_price"`
}

type StockInfo struct {
	TsCode     string         `db:"ts_code"`
	Symbol     string         `db:"symbol"`
	Name       string         `db:"name"`
	Area       string         `db:"area"`
	Industry   string         `db:"industry"`
	Fullname   string         `db:"fullname"`
	Enname     string         `db:"enname"`
	Cnspell    string         `db:"cnspell"`
	Market     string         `db:"market"`
	Exchange   string         `db:"exchange"`
	CurrType   string         `db:"curr_type"`
	ListStatus string         `db:"list_status"`
	ListDate   sql.NullString `db:"list_date"`
	DelistDate sql.NullString `db:"delist_date"`
	IsHs       string         `db:"is_hs"`
	ActName    sql.NullString `db:"act_name"`
	ActEntType sql.NullString `db:"act_ent_type"`
}

func loadConfig() (*Config, error) {
	file, err := os.Open("config.json")
	if err != nil {
		return nil, fmt.Errorf("打开配置文件失败: %v", err)
	}
	defer file.Close()

	var config Config
	decoder := json.NewDecoder(file)
	if err := decoder.Decode(&config); err != nil {
		return nil, fmt.Errorf("解析配置文件失败: %v", err)
	}

	return &config, nil
}

func getStockInfo(db *sqlx.DB, symbol string) (*StockInfo, error) {
	query := `
		SELECT ts_code, symbol, name, area, industry, fullname, enname, cnspell,
			   market, exchange, curr_type, list_status, list_date, delist_date,
			   is_hs, act_name, act_ent_type
		FROM stock_basic
		WHERE symbol = ?
		LIMIT 1
	`
	var stockInfo StockInfo
	err := db.Get(&stockInfo, query, symbol)
	if err != nil {
		return nil, err
	}
	return &stockInfo, nil
}

func insertStockPool(db *sqlx.DB, stockInfo *StockInfo) error {
	query := `
		INSERT INTO stock_pool (exchange, symbol, name, in_date, region, industry)
		VALUES (?, ?, ?, ?, ?, ?)
		ON DUPLICATE KEY UPDATE
		out_date = NULL
	`
	_, err := db.Exec(query,
		stockInfo.Exchange,
		stockInfo.Symbol,
		stockInfo.Name,
		time.Now().Format("2006-01-02"),
		stockInfo.Area,
		stockInfo.Industry,
	)
	return err
}

func runStockSelection(db *sqlx.DB) {
	// 获取最近一个月的股票数据
	oneMonthAgo := time.Now().AddDate(0, -1, 0)
	query := `
		SELECT symbol, exchange, datetime, ` + "`interval`" + `, volume, turnover, open_interest,
			   open_price, high_price, low_price, close_price
		FROM dbbardata
		WHERE datetime >= ?
		AND ` + "`interval`" + ` = 'd'
		AND (symbol LIKE '60%' OR symbol LIKE '00%')
		ORDER BY symbol, datetime DESC
	`

	var stocks []StockData
	err := db.Select(&stocks, query, oneMonthAgo)
	if err != nil {
		log.Printf("查询数据失败: %v", err)
		return
	}

	// 按股票代码分组处理数据
	stockMap := make(map[string][]StockData)
	for _, stock := range stocks {
		key := stock.Symbol
		stockMap[key] = append(stockMap[key], stock)
	}

	// 筛选符合条件的股票
	var qualifiedStocks []string
	for symbol, data := range stockMap {
		if isQualified(data) {
			qualifiedStocks = append(qualifiedStocks, symbol)
		}
	}

	// 输出结果并写入股票池
	log.Printf("开始执行选股，时间：%s", time.Now().Format("2006-01-02 15:04:05"))
	log.Printf("符合条件的股票数量：%d", len(qualifiedStocks))

	for _, symbol := range qualifiedStocks {
		log.Printf("处理股票：%s", symbol)

		// 获取股票详细信息
		stockInfo, err := getStockInfo(db, symbol)
		if err != nil {
			log.Printf("获取股票 %s 信息失败: %v", symbol, err)
			continue
		}

		// 写入股票池
		err = insertStockPool(db, stockInfo)
		if err != nil {
			log.Printf("写入股票池失败 %s: %v", symbol, err)
			continue
		}
	}
	log.Printf("选股完成，时间：%s", time.Now().Format("2006-01-02 15:04:05"))
}

func main() {
	// 加载配置文件
	config, err := loadConfig()
	if err != nil {
		log.Fatalf("加载配置失败: %v", err)
	}

	// 构建数据库连接字符串
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?parseTime=true",
		config.Database.Username,
		config.Database.Password,
		config.Database.Host,
		config.Database.Port,
		config.Database.DBName,
	)

	// 数据库连接配置
	db, err := sqlx.Connect("mysql", dsn)
	if err != nil {
		log.Fatalf("数据库连接失败: %v", err)
	}
	defer db.Close()

	// 程序启动时立即执行一次选股
	log.Printf("程序启动，立即执行一次选股")
	runStockSelection(db)

	// 设置定时执行
	for {
		now := time.Now()
		nextRun := time.Date(now.Year(), now.Month(), now.Day(), 18, 0, 0, 0, now.Location())
		if now.After(nextRun) {
			nextRun = nextRun.Add(24 * time.Hour)
		}

		// 计算等待时间
		waitTime := nextRun.Sub(now)
		log.Printf("下次执行时间：%s，等待时间：%v", nextRun.Format("2006-01-02 15:04:05"), waitTime)

		// 等待到执行时间
		time.Sleep(waitTime)

		// 执行选股
		runStockSelection(db)
	}
}

// 判断股票是否满足条件
func isQualified(data []StockData) bool {
	if len(data) < 10 {
		return false
	}

	// 查找最近10个交易日内的涨停
	var lastLimitUpIndex int = -1
	for i := 0; i < min(10, len(data)); i++ {
		if isLimitUp(data[i]) {
			lastLimitUpIndex = i
		}
	}

	if lastLimitUpIndex == -1 {
		return false
	}

	// 检查最后一次涨停后的所有收盘价是否都大于涨停价
	limitUpPrice := data[lastLimitUpIndex].ClosePrice
	for i := 0; i < lastLimitUpIndex; i++ {
		if data[i].ClosePrice <= limitUpPrice {
			return false
		}
	}

	return true
}

// 判断是否涨停
func isLimitUp(data StockData) bool {
	// 涨停判断逻辑：当日收盘价等于当日最高价，且涨幅大于等于9.9%
	// 由于没有前收盘价，我们使用前一天的收盘价作为参考
	limitUpRate := 0.099
	return data.ClosePrice == data.HighPrice && data.ClosePrice >= data.OpenPrice*(1+limitUpRate)
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
